import tensorflow as tf
from tensorflow.python.framework import graph_util

output_nodes = ['outputs']

def load_pb(path='model.pb'):
    with tf.gfile.GFile(path, 'rb') as fgraph:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(fgraph.read())

        return graph_def


def combined_graph():
    with tf.Graph().as_default() as g_combine:
        with tf.Session(graph=g_combine) as sess:
            graph_a = load_pb('graph_a.pb')
            graph_b = load_pb('graph_b.pb')

            model_a = tf.import_graph_def(graph_a, name='')
            model_b = tf.import_graph_def(graph_b, name='')

            g_combine_def = graph_util.convert_variables_to_constants(
                           sess=sess,
                           input_graph_def=sess.graph_def,
                           output_node_names=output_nodes)
            tf.train.write_graph(g_combine_def, './', 'model_combine.pb', as_text=False)
